import cv2
import numpy as np
import os
from control import *

LABELS = open("classes.names").read().strip().split("\n")
net = cv2.dnn.readNetFromDarknet(r'D:\xunlian\yolo_opencv\222\yolov4-tiny.cfg',
                                 r'D:\xunlian\yolo_opencv\222\trained_models\yolov4-tiny_final.weights')
layer = net.getUnconnectedOutLayersNames()

# 定义要处理的图片文件夹路径
image_folder = r'D:\lingsih\custom_training~1\home\hdkj\test\aa\custom_training\JPEGImages'  # 替换为你的图片文件夹路径

# 获取文件夹中所有的图片文件名列表
image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]

for image_file in image_files:
    frame = cv2.imread(os.path.join(image_folder, image_file))
    (H, W) = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)
    net.setInput(blob)
    layerOutputs = net.forward(layer)

    boxes = []
    confidences = []
    classIDs = []

    for output in layerOutputs:
        for detection in output:
            scores = detection[5:]
            classID = np.argmax(scores)
            confidence = scores[classID]
            if confidence > 0.1:  # 只有当置信度大于0.1时才认为是有效的检测结果
                box = detection[0:4] * np.array([W, H, W, H])
                (centerX, centerY, width, height) = box.astype("int")
                x = int(centerX - (width / 2))
                y = int(centerY - (height / 2))
                boxes.append([x, y, int(width), int(height)])
                confidences.append(float(confidence))
                classIDs.append(classID)

    idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.1, 0.9)
    if len(idxs) > 0:
        for i in idxs.flatten():
            (x, y) = (boxes[i][0], boxes[i][1])
            (w, h) = (boxes[i][2], boxes[i][3])
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2, lineType=cv2.LINE_AA)
            text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i])
            cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1, lineType=cv2.LINE_AA)

    frame = cv2.resize(frame, (960, 720))

    # 显示或保存处理后的图像
    # 如果想显示图像，请取消下两行的注释
    # cv2.imshow('frame', frame)
    # cv2.waitKey(0)

    # 如果想保存处理后的图像，请取消下一行的注释并确保输出路径正确
    cv2.imwrite(os.path.join('./output_folder', image_file), frame)

cv2.destroyAllWindows()